'''
Created on Jan 16, 2012

@author: oabalbin
'''
import numpy as np
from collections import defaultdict

def read_file(file, nsamples,pedigree):
    '''
    Y       3280031 0       G/T     0/1     5
    '''
    ifile = open(file)
    snps = defaultdict(list)
    snpids= defaultdict(list)
    
    for i,l in enumerate(ifile):
        f=l.strip('\n').split('\t')
        id='M_'+f[1]+'_'+f[2]
        snps[id].append((f[3],f[5]))
            
    ltable = np.zeros((nsamples, len(snps)*2 ),dtype='int')
    
    y={'A':1,'C':2,'G':3,'T':4}
    
    markers=[]
    of1 = open(file+'_linkage.info','w')
    snps_sort_keys = snps.keys()
    snps_sort_keys.sort()
    
    for j,p in enumerate(snps_sort_keys):
    #for j,p in enumerate(snps):
        sps = snps[p]
        # get the marker name
        
        name=p.split('_')[0]+str(j)
        markers.append(name)
        ol = [name, p.split('_')[1]]
        of1.write(",".join( list(map(str,ol)) ).replace(',','\t')+'\n')
        for s in sps:
            if int(s[1])==0: # 0 means all the samples share that snp
                g = s[0] # samples are being numbered starting at 1
                n,m = y[g.split('/')[0]], y[g.split('/')[1]]
                if m!=n:
                    n,m=1,2
                else:
                    n,m=2,2
                                    
                ltable[:,j*2],ltable[:,j*2+1] = n,m
            else:
                g,i=s[0], int(s[1])-1 # samples are being numbered starting at 1
                n,m = y[g.split('/')[0]], y[g.split('/')[1]]
                if m!=n:
                    n,m=1,2
                else:
                    n,m=2,2

                ltable[i,j*2],ltable[i,j*2+1] = n,m
    
    of = open(file+'_linkage.ped','w')
    
    #of.write(",".join( markers ).replace(',','\t')+'\n')    
    
    for i in range(ltable.shape[0]):
        ol =  ltable[i,:]
        ol1= pedigree[i,:]
        of.write(",".join( list(map(str,ol1)) ).replace(',',' ')+' '+",".join( list(map(str,ol)) ).replace(',',' ')+'\n')
    
    of.close()
    of1.close()
    return snps, snpids


file = '/exds/users/oabalbin/projects/fmpn/data/gatk_calls/haploview_chr1_all.txt'


#Pedigree    ID    Father    Mother    Sex    TRAIT.A

pedigree = np.array([[1,    1,    0,    0,    1,    1],
[1,    2,    0,    0,    2,    2],
[1,    3,    1,    2,    1,    0],
[1,    4,    1,    2,    1,    1],
[1,    5,    1,    2,    1,    0],
[1,    6,    0,    0,    1,    1],
[1,    7,    0,    0,    2,    2],
[1,    8,    6,    7,    1,    0],
[1,    9,    6,    7,    1,    0],
[1,    10,    0,    0,    1,    2]],dtype='int')

read_file(file,10,pedigree)



